%%--------------------------------------------------------------------
%% Copyright (c) 2025 EMQ Technologies Co., Ltd. All Rights Reserved.
%%--------------------------------------------------------------------

-module(emqx_mq_registry).

-moduledoc """
The module contains the registry of Message Queues.
""".

-include("emqx_mq_internal.hrl").

-export([
    create_tables/0,
    create/1,
    find/1,
    is_present/1,
    match/1,
    delete/1,
    update/2,
    list/0,
    list/2
]).

-dialyzer(no_improper_lists).

%% Only for testing/debugging.
-export([
    delete_all/0
]).

-define(MQ_REGISTRY_INDEX_TAB, emqx_mq_registry_index).
-define(MQ_REGISTRY_SHARD, emqx_mq_registry_shard).

-record(?MQ_REGISTRY_INDEX_TAB, {
    key :: emqx_topic_index:key(nil()) | '_',
    id :: emqx_mq_types:mqid() | '_',
    is_lastvalue :: boolean() | '_',
    key_expr :: emqx_variform:compiled() | undefined | '_',
    extra = #{} :: map() | '_'
}).

-type cursor() :: binary() | undefined.

%%--------------------------------------------------------------------
%% API
%%--------------------------------------------------------------------

-spec create_tables() -> [atom()].
create_tables() ->
    ok = mria:create_table(?MQ_REGISTRY_INDEX_TAB, [
        {type, ordered_set},
        {rlog_shard, ?MQ_REGISTRY_SHARD},
        {storage, disc_copies},
        {record_name, ?MQ_REGISTRY_INDEX_TAB},
        {attributes, record_info(fields, ?MQ_REGISTRY_INDEX_TAB)},
        {storage_properties, [
            {ets, [
                {read_concurrency, true},
                {write_concurrency, true}
            ]}
        ]}
    ]),
    [?MQ_REGISTRY_INDEX_TAB].

-doc """
Create a new MQ.
""".
-spec create(emqx_mq_types:mq()) ->
    {ok, emqx_mq_types:mq()}
    | {error, queue_exists}
    | {error, max_queue_count_reached}
    | {error, term()}.
create(#{topic_filter := TopicFilter, is_lastvalue := IsLastValue, limits := Limits} = MQ0) when
    (not IsLastValue) orelse (IsLastValue andalso is_map(map_get(key_expression, MQ0)))
->
    Key = make_key(TopicFilter),
    Id = emqx_guid:gen(),
    maybe
        ok ?= validate_max_queue_count(),
        {atomic, ok} ?=
            mria:transaction(?MQ_REGISTRY_SHARD, fun() ->
                case mnesia:read(?MQ_REGISTRY_INDEX_TAB, Key, write) of
                    [] ->
                        ok = mnesia:write(#?MQ_REGISTRY_INDEX_TAB{
                            key = Key,
                            id = Id,
                            is_lastvalue = IsLastValue,
                            key_expr = maps:get(key_expression, MQ0, undefined),
                            extra = #{
                                limits => Limits
                            }
                        }),
                        ok;
                    [_] ->
                        {error, queue_exists}
                end
            end),
        MQ = MQ0#{id => Id},
        try emqx_mq_state_storage:create_mq_state(MQ) of
            {ok, _} ->
                {ok, MQ};
            {error, Reason} ->
                ?tp(error, mq_registry_create_mq_state_error, #{
                    mq => MQ,
                    reason => Reason
                }),
                mria:dirty_delete(?MQ_REGISTRY_INDEX_TAB, Key),
                {error, Reason}
        catch
            Class:Reason ->
                ?tp(error, mq_registry_create_mq_state_error, #{
                    mq => MQ,
                    class => Class,
                    reason => Reason
                }),
                mria:dirty_delete(?MQ_REGISTRY_INDEX_TAB, Key),
                {error, Reason}
        end
    else
        {atomic, {error, _} = Error} ->
            Error;
        {aborted, ErrReason} ->
            {error, ErrReason};
        {error, _} = Error ->
            Error
    end.

-doc """
Find all MQs matching the given concrete topic.
""".
-spec match(emqx_types:topic()) -> [emqx_mq_types:mq_handle()].
match(Topic) ->
    Keys = emqx_topic_index:matches(Topic, ?MQ_REGISTRY_INDEX_TAB, []),
    lists:flatmap(
        fun(Key) ->
            case mnesia:dirty_read(?MQ_REGISTRY_INDEX_TAB, Key) of
                [] ->
                    [];
                [#?MQ_REGISTRY_INDEX_TAB{} = Rec] ->
                    [record_to_mq_handle(Rec)]
            end
        end,
        Keys
    ).

-doc """
Find the MQ by its topic filter.
""".
-spec find(emqx_mq_types:mq_topic()) -> {ok, emqx_mq_types:mq()} | not_found.
find(TopicFilter) ->
    ?tp_debug(mq_registry_find, #{topic_filter => TopicFilter}),
    Key = make_key(TopicFilter),
    case mnesia:dirty_read(?MQ_REGISTRY_INDEX_TAB, Key) of
        [] ->
            not_found;
        [#?MQ_REGISTRY_INDEX_TAB{id = Id}] ->
            emqx_mq_state_storage:find_mq(Id)
    end.

-doc """
Check if the MQ exists by its topic filter.
""".
-spec is_present(emqx_mq_types:mq_topic()) -> boolean().
is_present(TopicFilter) ->
    Key = make_key(TopicFilter),
    case mnesia:dirty_read(?MQ_REGISTRY_INDEX_TAB, Key) of
        [] ->
            false;
        [#?MQ_REGISTRY_INDEX_TAB{}] ->
            true
    end.

-doc """
Delete the MQ by its topic filter.
""".
-spec delete(emqx_mq_types:mq_topic()) -> ok | not_found | {error, term()}.
delete(TopicFilter) ->
    ?tp_debug(mq_registry_delete, #{topic_filter => TopicFilter}),
    Key = make_key(TopicFilter),
    case mnesia:dirty_read(?MQ_REGISTRY_INDEX_TAB, Key) of
        [] ->
            not_found;
        [#?MQ_REGISTRY_INDEX_TAB{} = Rec] ->
            #{id := Id} = MQHandle = record_to_mq_handle(Rec),
            ok = mria:dirty_delete_object(Rec),
            case emqx_mq_consumer:find(Id) of
                {ok, ConsumerRef} ->
                    ok = emqx_mq_consumer:stop(ConsumerRef);
                not_found ->
                    ok
            end,
            maybe
                ok ?= emqx_mq_message_db:drop(MQHandle),
                ok ?= emqx_mq_state_storage:destroy_consumer_state(MQHandle),
                ok ?= emqx_mq_state_storage:destroy_mq_state(MQHandle)
            end
    end.

-doc """
Delete all MQs. Only for testing/maintenance.
""".
-spec delete_all() -> ok.
delete_all() ->
    _ = mria:clear_table(?MQ_REGISTRY_INDEX_TAB),
    _ = emqx_mq_state_storage:delete_all(),
    ok.

-doc """
Update the MQ by its topic filter.
* `is_lastvalue` flag cannot be updated.
* limited regular queues cannot be updated to unlimited regular queues and vice versa.
""".
-spec update(emqx_mq_types:mq_topic(), map()) ->
    {ok, emqx_mq_types:mq()}
    | not_found
    | {error, is_lastvalue_not_allowed_to_be_updated}
    | {error, term()}.
update(TopicFilter, #{is_lastvalue := _IsLastValue} = UpdateFields0) ->
    Key = make_key(TopicFilter),
    UpdateFields = maps:without([topic_filter, id], UpdateFields0),
    case mnesia:dirty_read(?MQ_REGISTRY_INDEX_TAB, Key) of
        [] ->
            not_found;
        [#?MQ_REGISTRY_INDEX_TAB{} = Rec] ->
            #{id := Id} = MQHandle = record_to_mq_handle(Rec),
            IsLastvalueOld = emqx_mq_prop:is_lastvalue(MQHandle),
            IsLastvalueNew = emqx_mq_prop:is_lastvalue(UpdateFields),
            IsLimitedOld = emqx_mq_prop:is_limited(MQHandle),
            IsLimitedNew = emqx_mq_prop:is_limited(UpdateFields),
            NeedUpdateIndex = need_update_index(MQHandle, UpdateFields),
            case UpdateFields of
                _ when IsLastvalueOld =/= IsLastvalueNew ->
                    {error, is_lastvalue_not_allowed_to_be_updated};
                _ when (not IsLastvalueNew) andalso (IsLimitedOld =/= IsLimitedNew) ->
                    {error, limit_presence_cannot_be_updated_for_regular_queues};
                _ when NeedUpdateIndex ->
                    case update_index(Key, Id, UpdateFields) of
                        ok ->
                            emqx_mq_state_storage:update_mq_state(Id, UpdateFields);
                        not_found ->
                            not_found
                    end;
                _ ->
                    emqx_mq_state_storage:update_mq_state(Id, UpdateFields)
            end
    end.

-doc """
List all MQs.
""".
-spec list() -> emqx_utils_stream:stream(emqx_mq_types:mq()).
list() ->
    mq_record_stream_to_queues(mq_record_stream()).

-doc """
List at most `Limit` MQs starting from `Cursor` position.
""".
-spec list(cursor(), non_neg_integer()) -> {[emqx_mq_types:mq()], cursor()}.
list(Cursor, Limit) when Limit >= 1 ->
    MQs0 = emqx_utils_stream:consume(
        emqx_utils_stream:limit_length(
            Limit + 1,
            mq_record_stream_to_queues(mq_record_stream(Cursor))
        )
    ),
    case length(MQs0) < Limit + 1 of
        true ->
            {MQs0, undefined};
        false ->
            MQs = lists:sublist(MQs0, Limit),
            #{topic_filter := TopicFilter} = lists:last(MQs),
            NewCursor = TopicFilter,
            {MQs, NewCursor}
    end.

%%--------------------------------------------------------------------
%% Internal functions
%%--------------------------------------------------------------------

need_update_index(
    #{is_lastvalue := true, key_expression := OldKeyExpr} = _MQHandle,
    #{key_expression := NewKeyExpr} = _UpdateFields
) when NewKeyExpr =/= OldKeyExpr ->
    true;
need_update_index(#{limits := OldLimits} = _MQHandle, #{limits := NewLimits} = _UpdateFields) when
    OldLimits =/= NewLimits
->
    true;
need_update_index(_, _) ->
    false.

mq_record_stream() ->
    mq_record_stream(undefined).

mq_record_stream(Cursor) ->
    Stream = mq_ets_record_stream(key_from_cursor(Cursor)),
    emqx_utils_stream:chainmap(
        fun(L) -> L end,
        Stream
    ).

key_from_cursor(undefined) ->
    undefined;
key_from_cursor(Cursor) ->
    make_key(Cursor).

mq_ets_record_stream(Key) ->
    fun() ->
        case next_key(Key) of
            '$end_of_table' ->
                [];
            NextKey ->
                [ets:lookup(?MQ_REGISTRY_INDEX_TAB, NextKey) | mq_ets_record_stream(NextKey)]
        end
    end.

next_key(undefined) ->
    ets:first(?MQ_REGISTRY_INDEX_TAB);
next_key(Key) ->
    ets:next(?MQ_REGISTRY_INDEX_TAB, Key).

mq_record_stream_to_queues(Stream) ->
    emqx_utils_stream:chainmap(
        fun(#?MQ_REGISTRY_INDEX_TAB{id = Id}) ->
            case emqx_mq_state_storage:find_mq(Id) of
                {ok, MQ} ->
                    [MQ];
                not_found ->
                    []
            end
        end,
        Stream
    ).

make_key(TopicFilter) ->
    emqx_topic_index:make_key(TopicFilter, []).

update_index(Key, Id, UpdateFields) ->
    {atomic, Result} = mria:transaction(?MQ_REGISTRY_SHARD, fun() ->
        case mnesia:read(?MQ_REGISTRY_INDEX_TAB, Key, write) of
            [] ->
                not_found;
            [#?MQ_REGISTRY_INDEX_TAB{id = Id, extra = Extra} = Rec] ->
                NewKeyExpr = maps:get(key_expression, UpdateFields, undefined),
                NewLimits = maps:get(limits, UpdateFields, ?DEFAULT_MQ_LIMITS),
                mnesia:write(Rec#?MQ_REGISTRY_INDEX_TAB{
                    key_expr = NewKeyExpr, extra = Extra#{limits => NewLimits}
                }),
                ok
        end
    end),
    Result.

record_to_mq_handle(#?MQ_REGISTRY_INDEX_TAB{
    key = Key, id = Id, key_expr = KeyExpr, is_lastvalue = IsLastValue, extra = Extra
}) ->
    #{
        id => Id,
        topic_filter => emqx_topic_index:get_topic(Key),
        is_lastvalue => IsLastValue,
        key_expression => KeyExpr,
        limits => maps:get(limits, Extra, ?DEFAULT_MQ_LIMITS)
    }.

queue_count() ->
    mnesia:table_info(?MQ_REGISTRY_INDEX_TAB, size).

validate_max_queue_count() ->
    case queue_count() >= emqx_mq_config:max_queue_count() of
        true ->
            {error, max_queue_count_reached};
        false ->
            ok
    end.
